#pragma once
#include "LightSource.hpp"
#include "Photon.hpp"
#include "Media.hpp"
#include "Target.hpp"
#include "Sensor.hpp"
#include "Params.hpp"
#include "constants.hpp"
#include <fstream>
#include <vector>
#include <omp.h>
#include <cmath>


#define min(a, b) ((a) < (b) ? (a) : (b))
#define max(a, b) ((a) > (b) ? (a) : (b))
#define SIGN(x) ((x)>=0 ? 1:-1)

class Simulation
{
public:
    Simulation();
    Simulation(double a, double b, double zt, int_fast64_t NPHOTON);
    Simulation(std::vector<LightSource> vecLightSources, Media media, 
            Target target, Sensor sensor, Params params);
    Simulation(Simulation& simulation);
    Simulation& operator=(const Simulation& simulation);
    void start();
    void createPhoton(Photon& photon, int_fast64_t iphoton);
    void propagatePhoton(Photon& photon);

    bool target(double x,double y);

    // Random number generation related
    void InitRand();
    double UnifClosed();   // [0, 1]
    double UnifOpen();     // ]0 ,1[
    double UnifHalfDown(); // ]0, 1]
    double UnifHalfUp();   // [0, 1[

public:
    // std::vector<LightSource> vecLightSources;
    // Media media;
    // Target target;
    // Sensor sensor;
    // Params params;

    mt_rng rng;                         // random number genetor
    unsigned long seed;

    double a,b,c,g;                     // absorption, scattering, attenuation coefficients and asymmetry factor
    double g2;                          // for fast compution
    int_fast64_t NPHOTON;               // number of photons to be simulated
    double THRESHOLD;                   // threshold to terminate tracking photon
    double CHANCE;                      // chance for photon to be revival

    // constant parameters
    double EPS;                         // minimum value to avoid division by zero

    // // enum parameters
    // int enumMediaType;                  // enum variable for media type
    // int enumMethodType;                 // enum variable for simulation method
    // int sourceType;                     // light source type
    // bool isPolarized;                   // polarization?
    // bool isPhase;                       // phase variant?
    // bool isTime;                        // time variant?
    // Point3d sourceCentralPosition;      // central position of light source
    // Vec3d sourceCentralDirection;       // central direction of light source
    double waveLength;                  // wavelength of light
    // double photonsPerJoule;             // sampling rate, photon number / pulse energy (joule)

    // for time variant source
    // double timeScale;                   // time step of input time-intensity array
    double tp;                          // pulse time
    // std::function<double(double)> pulseShape;   // pulse shape function

    // spatial parameters
    double divengencyAngle;             // divengency angle of light source, deg
    // double sourceArea;                  // area of light source area, for gaussian and area type

    // sensor
    // int pixelX;
    // int pixelY;
    // double element;
    double fov;
    double tg;
    // Point3d sensorCenter;
    // Vec3d sensorNormal;

    // target
    // Point3d targetCenter;
    // Vec3d targetNormal;
    // cv::Mat targetImage;
    // double sizeScale;       // real size (m) / pixel size (pixel)

    // OpeMP & MPI related variables
    int threadcount;
    int rank, nodecount;
    int totalthreads;

    // temp use for 1st edition
    double laserX,laserY,laserZ;
    double sigma,wd,d,zr;
    double zt;

    // for record data
    int_fast64_t nn;
    std::vector<Data> photonData;
};


Simulation::Simulation()
{
    seed = 5489UL;
    a = 0.5;
    b = 0.5;
    c = a+b;
    g = 0.9;
    g2 = g*g;

    NPHOTON = 1000000;
    THRESHOLD = 1.0e-4;
    CHANCE = 0.1;
    EPS = 1.0e-12;

    waveLength = 532e-9;
    tp = 5.0e-9;
    divengencyAngle = 5.0 / 180 * PI / 2.0;
    // sourceArea = 0;

    fov = 60;
    tg = tp;
    rank = 0;
    threadcount = nodecount = totalthreads = 1;
    laserX = 0;
    laserY = 0;
    laserZ = 0;

    sigma = 1.8;
    zr = PI * sigma * sigma / waveLength;
    d = sqrt(1.0 / ((divengencyAngle / sigma) * (divengencyAngle / sigma) - 1.0 / (zr * zr)));
    wd = sigma * sqrt(1.0 + (d / zr) * (d / zr));


    zt = 3.0;

    nn = (int_fast64_t)((log10(CHANCE)/(log10(b)-log10(c))+1/0.9)*2*(double)NPHOTON);
    photonData.reserve(nn);
}


Simulation::Simulation(Simulation& simulation)
{
    this->seed = simulation.seed;
    this->a = simulation.a;
    this->b = simulation.b;
    this->c = simulation.c;
    this->g = simulation.g;
    this->g2 = simulation.g2;
    this->NPHOTON = simulation.NPHOTON;
    this->THRESHOLD = simulation.THRESHOLD;
    this->CHANCE = simulation.CHANCE;
    this->EPS = simulation.EPS;
    this->waveLength = simulation.waveLength;
    this->tp = simulation.tp;
    this->divengencyAngle = simulation.divengencyAngle;
    this->fov = simulation.fov;
    this->tg = simulation.tg;
    this->threadcount = simulation.threadcount;
    this->rank = simulation.rank;
    this->nodecount = simulation.nodecount;
    this->totalthreads = simulation.totalthreads;
    this->laserX = simulation.laserX;
    this->laserY = simulation.laserY;
    this->laserZ = simulation.laserZ;
    this->sigma = simulation.sigma;
    this->wd = simulation.wd;
    this->d = simulation.d;
    this->zr = simulation.zr;
    this->zt = simulation.zt;
    this->nn = simulation.nn;
    // this->photonData.reserve(nn);
}


Simulation::Simulation(double a, double b, double zt, int_fast64_t NPHOTON)
{
    seed = 5489UL;
    this->a = a;
    this->b = b;
    c = a+b;
    g = 0.9;
    g2 = g*g;

    this->NPHOTON = NPHOTON;
    THRESHOLD = 1.0e-4;
    CHANCE = 0.1;
    EPS = 1.0e-12;

    waveLength = 532e-9;
    tp = 5.0e-9;
    divengencyAngle = 5.0 / 180 * PI / 2.0;
    // sourceArea = 0;

    fov = 60;
    tg = tp;
    rank = 0;
    threadcount = nodecount = totalthreads = 1;
    laserX = 0;
    laserY = 0;
    laserZ = 0;

    sigma = 1.8;
    zr = PI * sigma * sigma / waveLength;
    d = sqrt(1.0 / ((divengencyAngle / sigma) * (divengencyAngle / sigma) - 1.0 / (zr * zr)));
    // std::cout<<zr<<std::endl;
    wd = sigma * sqrt(1.0 + (d / zr) * (d / zr));


    this->zt = zt;

    nn = (int_fast64_t)((log10(CHANCE)/(log10(b)-log10(c))+1/0.9)*2*(double)NPHOTON);
    photonData.reserve(nn);
}


Simulation& Simulation::operator=(const Simulation& simulation)
{
    this->seed = simulation.seed;
    this->a = simulation.a;
    this->b = simulation.b;
    this->c = simulation.c;
    this->g = simulation.g;
    this->g2 = simulation.g2;
    this->NPHOTON = simulation.NPHOTON;
    this->THRESHOLD = simulation.THRESHOLD;
    this->CHANCE = simulation.CHANCE;
    this->EPS = simulation.EPS;
    this->waveLength = simulation.waveLength;
    this->tp = simulation.tp;
    this->divengencyAngle = simulation.divengencyAngle;
    this->fov = simulation.fov;
    this->tg = simulation.tg;
    this->threadcount = simulation.threadcount;
    this->rank = simulation.rank;
    this->nodecount = simulation.nodecount;
    this->totalthreads = simulation.totalthreads;
    this->laserX = simulation.laserX;
    this->laserY = simulation.laserY;
    this->laserZ = simulation.laserZ;
    this->sigma = simulation.sigma;
    this->wd = simulation.wd;
    this->d = simulation.d;
    this->zr = simulation.zr;
    this->zt = simulation.zt;
    this->nn = simulation.nn;
    // this->photonData.reserve(nn);
    return *this;
}


void Simulation::start()
{
    std::cout<<"starting simulation..."<<std::endl;
    threadcount = omp_get_max_threads();
    std::cout<<"total thread number: "<<threadcount<<std::endl;
    InitRand();
    std::cout<<"random number generator initialized..."<<std::endl;

    int_fast64_t ii, jj, nthread = omp_get_max_threads();
    int_fast64_t *ticks = new int_fast64_t[(int)nthread];

    Simulation *MCS = new Simulation[(int)nthread];
    bool abort_computation = false;

    // initialize threads
    for (ii = 0; ii < nthread; ii++)
    {
        std::cout<<"initializing thread "<<ii<<"..."<<std::endl;
        MCS[ii] = *this;
        MCS[ii].NPHOTON = NPHOTON / nthread;
        MCS[ii].nn = (int_fast64_t)((log10(CHANCE)/(log10(b)-log10(c))+1/0.9)*2*(double)(NPHOTON/nthread));
        MCS[ii].photonData.reserve(nn);
        MCS[ii].seed = (unsigned long) (MCS[ii].seed * totalthreads + ii);
        MCS[ii].InitRand();
        ticks[ii] = 0;
        std::cout<<"thread "<<ii<<" initialized..."<<std::endl;
    }

    std::cout<<"all threads initialized..."<<std::endl;

    // [AL] if remainder of nphoton / nthread is non-zero, total photon count is not the same as Nphoton
    // therefore add the remaining photons to the last thread.
    int_fast64_t realnphot = 0;
    for (ii = 0; ii < nthread; ii++)
        realnphot += MCS[ii].NPHOTON;
    MCS[nthread - 1].NPHOTON += NPHOTON - realnphot;

    std::cout<<"simulation on all threads prepared..."<<std::endl;

    #define TICK_VAL 10000

    // Compute Monte Carlo on each thread seperately
    #pragma omp parallel
    {
        int_fast64_t iphoton, thread = omp_get_thread_num();
        std::cout<<"thread "<<thread<<" started..."<<std::endl;
        Photon phot;    // define photon
        int_fast64_t photonInitial = NPHOTON/nthread*thread;
        for (iphoton = photonInitial; iphoton < photonInitial+MCS[thread].NPHOTON; iphoton++)
        {
            ticks[thread] = iphoton-photonInitial+1;
            if (ticks[thread] % TICK_VAL == 0)
            {
                #pragma omp critical
                if (thread == 0)
                {
                    int_fast64_t jj, csum = 0;
                    {
                        for (jj = 0; jj < nthread; jj++)
                        {
                        csum += ticks[jj];
                        }
                        std::cout<<csum<<std::endl;
                        std::cout<<"simulation progress: "<<(100*(double)csum/(double)NPHOTON)<<"% ..."<<std::endl;
                    }
                }
            }
            MCS[thread].createPhoton(phot, iphoton);
            MCS[thread].propagatePhoton(phot);
        }
    }
    #pragma omp barrier

    std::cout<<"simulation finished..."<<std::endl;
    // Sum up the results to first instance and delete MCS
    NPHOTON = 0;
    std::ofstream fout;
    fout.open("a.txt",std::ios::out);
    if(fout.is_open())
    {
        std::cout<<"now start writing files..."<<std::endl;
        for (jj = 0; jj < nthread; jj++)
        {
            NPHOTON += MCS[jj].NPHOTON;
            for (auto i : MCS[jj].photonData)
            {
                fout<<i.x<<"\t"<<i.y<<"\t"<<i.z<<"\t"<<i.t<<"\t"<<i.W<<std::endl;
            }
            std::cout<<NPHOTON<<"photons saved..."<<std::endl;
        }
        std::cout<<"all photons saved..."<<std::endl;
    }
    else std::cout<<"failed to open file."<<std::endl;
    
    //  delete[] ticks;
    delete[] ticks;
    delete[] MCS;

    std::cout<<"program finished..."<<std::endl;
}


void Simulation::createPhoton(Photon& photon, int_fast64_t iphoton)
{
    // initialize photon status
    double r11 = UnifHalfDown();        // radial random number for xs, ys
    double r21 = UnifHalfDown();        // radial random number for xf, yf
    double r12 = UnifClosed();          // angular random number for xs, ys
    double r22 = UnifClosed();          // angular random number for xf, yf
	double xs = wd * sqrt(-log(r11)) * cos(2.0 * PI * r12);
	double ys = wd * sqrt(-log(r11)) * sin(2.0 * PI * r12);
	double xf = sigma * sqrt(-log(r21)) * cos(2.0 * PI * r22);
	double yf = sigma * sqrt(-log(r21)) * sin(2.0 * PI * r22);
	double tmp = sqrt((xs - xf) * (xs - xf) + (ys - yf) * (ys - yf) + d * d);
	photon.ux = (xs - xf) / tmp;
	photon.uy = (ys - yf) / tmp;
	photon.uz = d / tmp;
    // std::cout<<wd<<std::endl;
	photon.x = laserX + xs;
	photon.y = laserY + ys;
	photon.z = laserZ;
	photon.W = 1.0;
	photon.t = iphoton * tp / NPHOTON;
	this->zt = zt;
}


void Simulation::propagatePhoton(Photon& photon)
{
    while (1)
    {
        // move photon
        double l = -log(UnifOpen()) / c;
        photon.x += l * photon.ux;
        photon.y += l * photon.uy;
        photon.z += l * photon.uz;
        photon.t += l/speed;

        // target reflection?
        // record position
        if (photon.z >= zt)
        {
            double lt = l-(photon.z - zt)/photon.uz;
            // double temp = sqrt(1.0 - photon.uz * photon.uz) / photon.uz * (photon.z - zt);
            double x0 = photon.x - (photon.z - zt)/photon.uz * photon.ux;
            double y0 = photon.y - (photon.z - zt)/photon.uz * photon.uy;
            double t0 = photon.t - lt/speed;
            if (target(x0, y0))
            {
                Data data;
                data.x = x0;
                data.y = y0;
                data.z = zt;
                data.t = t0;
                data.W = photon.W;
                photonData.push_back(data);
                photon.z -= 2.0 * (photon.z - zt);
                photon.uz = -photon.uz;
                data.x = photon.x;
                data.y = photon.y;
                data.z = photon.z;
                data.t = photon.t;
                data.W = photon.W;
                photonData.push_back(data);
            }
            else break;
        }
        else if (photon.z<0) break;
        else
        {
            Data data;
            data.x = photon.x;
            data.y = photon.y;
            data.z = photon.z;
            data.t = photon.t;
            data.W = photon.W;
            photonData.push_back(data);
        }
        
        // scattering photon
        double costheta, sintheta, cospsi, sinpsi, uxx, uyy, uzz, temp;
        double rp = UnifClosed();   // random number for polar scattering
        double ra = UnifClosed();   // random number for azimuth scattering
        costheta = (1.0 + g2 - pow((1.0 - g2) / (1.0 - g * (1.0 - 2.0 * rp)), 2)) / (2.0 * g);
        sintheta = sqrt(1.0 - costheta * costheta);

        cospsi = cos(2.0 * PI * ra);
        sinpsi = sqrt(1.0 - cospsi * cospsi);

        if (1.0 - fabs(photon.uz) <= EPS)
        {
            uxx = sintheta * cospsi;
            uyy = sintheta * sinpsi;
            uzz = costheta * SIGN(photon.uz);
        }
        else
        {
            temp = sqrt(1.0 - photon.uz * photon.uz);
            uxx = sintheta * (photon.ux * photon.uz * cospsi - photon.uy * sinpsi) / temp + photon.ux * costheta;
            uyy = sintheta * (photon.uy * photon.uz * cospsi + photon.ux * sinpsi) / temp + photon.uy * costheta;
            uzz = -sintheta * cospsi * temp + photon.uz * costheta;
        }

        photon.ux = uxx;
        photon.uy = uyy;
        photon.uz = uzz;

        // set weight
        photon.W *= b / c;

        // set alive?
        if (photon.W < THRESHOLD)
        {
            if (UnifClosed() < CHANCE) photon.W /= CHANCE;
            else break;
        }
    }
}


// Initialize random number generator
void Simulation::InitRand()
{
  rng.Seed(seed);
}

// Draw random number on [0, 1]
double Simulation::UnifClosed()
{
  return (rng.drand_closed());
}

// Draw random number on ]0, 1[
double Simulation::UnifOpen()
{
  return (rng.drand_open());
}

// Draw random number ]0, 1]
double Simulation::UnifHalfDown()
{
  return (rng.drand_open_down());
}

// Draw random number [0, 1[
double Simulation::UnifHalfUp()
{
  return (rng.drand_open_up());
}


bool Simulation::target(double x,double y)
{
    return true;    // for expand object

    // for nonexpand object
    // return object(double x,double y);
}

